import grpc
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import dtypes
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc

from constant.constant_config import Data
from constant.constant_config import data_prefix
from constant.constant_config import max_click_history
from utils.common_util import get_app
from utils.redis_util import redis_template

app = get_app()


# 根据uid进行推荐
def recommend_entity_by_UID(uid):
    data = redis_template.get_list(data_prefix + str(uid))
    data = transform(data)
    # par:max_click_history
    uid2words, uid2entities = aggregate(data, max_click_history=max_click_history)
    data = generate_model_data(data, uid2words, uid2entities)
    channel = grpc.insecure_channel(app.config['MODEL_HOST'] + ':' + str(app.config['MODEL_PORT']))
    stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
    request = predict_pb2.PredictRequest()
    request.model_spec.name = app.config['MODEL_SPE_NAME']
    request.model_spec.signature_name = app.config['MODEL_SIGNATURE_NAME']
    request.inputs['clicked_words'].CopyFrom(
        tf.make_tensor_proto(data.clicked_words, shape=[1, data.clicked_words.shape[0], data.clicked_words.shape[1]],
                             dtype=dtypes.int32))
    request.inputs['clicked_entities'].CopyFrom(tf.make_tensor_proto(data.clicked_entities,
                                                                     shape=[1, data.clicked_entities.shape[0],
                                                                            data.clicked_entities.shape[1]],
                                                                     dtype=dtypes.int32))
    request.inputs['news_words'].CopyFrom(
        tf.make_tensor_proto(data.news_words, shape=[1, data.news_words.shape[0], data.news_words.shape[1]],
                             dtype=dtypes.int32))
    request.inputs['news_entities'].CopyFrom(
        tf.make_tensor_proto(data.news_entities, shape=[1, data.news_entities.shape[0], data.news_entities.shape[1]],
                             dtype=dtypes.int32))
    result_future = stub.Predict.future(request, app.config['MODEL_TIMEOUT'])
    result = result_future.result()
    scores = np.array(result.outputs[app.config['MODEL_OUTPUT_SCORE']].float_val)
    labels = np.array(result.outputs[app.config['MODEL_OUTPUT_LABELS']].float_val)
    return labels, scores


def aggregate(data, max_click_history):
    uid2words = dict()
    uid2entities = dict()
    pos_idx = np.where(data['label'] == 1)[0]
    words = data['news_words'][pos_idx]
    entities = data['news_entities'][pos_idx]
    indices = np.random.choice(list(range(0, pos_idx.shape[0])), size=max_click_history, replace=True)
    uid2words[data['user_id'][0]] = words[indices]
    uid2entities[data['user_id'][0]] = entities[indices]
    return uid2words, uid2entities


# 转换redis中的数据
def transform(datas):
    res = {'user_id': [], 'news_words': [], 'news_entities': [], 'label': []}
    for data in datas:
        dic = modify(str(data))
        res['user_id'].append(dic['user_id'])
        res['news_words'].append(dic['news_words'])
        res['news_entities'].append(dic['news_entities'])
        res['label'].append(dic['label'])
    res['user_id'] = np.array(res['user_id'])
    res['news_words'] = np.array(res['news_words'])
    res['news_entities'] = np.array(res['news_entities'])
    res['label'] = np.array(res['label'])
    return res


def generate_model_data(data, uid2words, uid2entities):
    data['clicked_words'] = uid2words[data['user_id'][0]]
    data['clicked_entities'] = uid2entities[data['user_id'][0]]
    data = Data(size=len(data),
                clicked_words=np.array(data['clicked_words']),
                clicked_entities=np.array(data['clicked_entities']),
                news_words=data['news_words'],
                news_entities=data['news_entities'],
                labels=data['label'])
    return data


# 修改数据
# example b'0\t1,2,3,0,0,0,0,0,0,0\t0,0,0,0,0,0,0,0,0,0\t0\n'
def modify(data):
    dic = {}
    data = data.strip().split('\\t')
    dic['user_id'] = int(data[0][2:])
    news_words = data[1].split(',')
    news_words = np.array([int(item) for item in news_words])
    dic['news_words'] = news_words
    news_entities = data[2].split(',')
    news_entities = np.array([int(item) for item in news_entities])
    dic['news_entities'] = news_entities
    dic['label'] = int(data[3][:-3])
    return dic


if __name__ == '__main__':
    recommend_entity_by_UID(0)
